/*
 * GeoEquilibrator2.cpp
 *
 *  Created on: 17 Oct 2011
 *      Author: allan
 */

#include "GeoEquilibrator2.h"

#include "Geochemistry/GeoState.h"
#include "Utils/SpeciesUtils.h"
#include "Utils/Assert.h"

const EquilibriumCondition SpeciesMassConditionHelper2(const string& species, const FunctionTPna& h, const Multiphase& multiphase);
const EquilibriumCondition SpeciesMassCondition2(const string& species, double mass, const Multiphase& multiphase);
const EquilibriumCondition SpeciesActivityCondition2(const string& species, double activity, const Multiphase& multiphase);
const EquilibriumCondition ElementMassCondition2(const string& element, double mass, const Multiphase& multiphase);
const EquilibriumCondition AcidityCondition2(double pH, const Multiphase& multiphase);
const EquilibriumCondition PartialPressureCondition2(const string& gas, double Pg, const Multiphase& multiphase);
const EquilibriumCondition ChargeBalanceCondition2(const Multiphase& multiphase);

GeoEquilibrator2::GeoEquilibrator2(const GeoSystem& geosystem) :
multiphase(geosystem.GetMultiphase()), eReactions(geosystem.GetEquilibriumReactions())
{
	eSolver.Initialiaze(multiphase, eReactions);
	eSolver.SetDecompositionAlgorithm(NewtonSolver::PartialPivotingLU);
	
	n = VectorXd::Constant(multiphase.GetNumSpecies(), 1.0E-7);
}

void GeoEquilibrator2::ImposeSpeciesMass(const string& species, double mass)
{
	conditions.push_back(SpeciesMassCondition2(species, mass, multiphase));
	
	SetCompositionGuess(species, mass);
}

void GeoEquilibrator2::ImposeSpeciesActivity(const string& species, double activity)
{
	conditions.push_back(SpeciesActivityCondition2(species, activity, multiphase));
	
	SetCompositionGuess(species, activity);
}

void GeoEquilibrator2::ImposeElementMass(const string& element, double mass)
{
	conditions.push_back(ElementMassCondition2(element, mass, multiphase));
}

void GeoEquilibrator2::ImposeAcidity(double pH)
{
	conditions.push_back(AcidityCondition2(pH, multiphase));
	
	SetCompositionGuess("H[+]", pow(10.0, -pH));
}

void GeoEquilibrator2::ImposePartialPressure(const string& gas, double Pg)
{
	conditions.push_back(PartialPressureCondition2(gas, Pg, multiphase));
	
	SetCompositionGuess(gas, Pg);
}

void GeoEquilibrator2::ImposeChargeBalance()
{
	conditions.push_back(ChargeBalanceCondition2(multiphase));
}

void GeoEquilibrator2::ResetConditions()
{
	conditions.clear();
	
	n = VectorXd::Constant(multiphase.GetNumSpecies(), 1.0E-7);
}

void GeoEquilibrator2::SetCompositionGuess(const string& species, double mass)
{
	n[multiphase[species]] = mass;
}

ConvergenceDiagnostics GeoEquilibrator2::Equilibrate(GeoState& state)
{
	EquilibriumConditions econditions(conditions);
	
	double T = state.GetTemperature();
	double P = state.GetPressure();
	
	ConvergenceDiagnostics diagnostics = eSolver.Solve(econditions, T, P, n);
	
	Assert(diagnostics.converged, "The equilibrium calculation did not converge.");
	
	state.SetComposition(n);
	
	return diagnostics;
}

ConvergenceDiagnostics GeoEquilibrator2::Equilibrate(GeoState& state, const VectorXd& nGuess)
{
	EquilibriumConditions econditions(conditions);
	
	double T = state.GetTemperature();
	double P = state.GetPressure();
	n = nGuess;
	
	ConvergenceDiagnostics diagnostics = eSolver.Solve(econditions, T, P, n);
	
	Assert(diagnostics.converged, "The equilibrium calculation did not converge.");
	
	state.SetComposition(n);
	
	return diagnostics;
}

const EquilibriumCondition SpeciesMassConditionHelper2(const string& species, const FunctionTPna& h, const Multiphase& multiphase)
{
	// The number of species in the system
	const unsigned N = multiphase.GetNumSpecies();
	
	// The equilibrium condition instance
	EquilibriumCondition eCondition(N);
	
	// Set the equilibrium condition function (h)
	eCondition.hj = h;
	
	// Get the index of the given (species) among all the species in the system
	const unsigned i = multiphase[species];
	
	// If the given (species) is a primary species, set the alphaRow vector of (eCondition)
	if(i < N) eCondition.alpha_row[i] = 1.0;
	
	return eCondition;
}

const EquilibriumCondition SpeciesMassCondition2(const string& species, double mass, const Multiphase& multiphase)
{
	return SpeciesMassConditionHelper2(species, ConstantTPna(mass), multiphase);
}

const EquilibriumCondition SpeciesActivityCondition2(const string& species, double activity, const Multiphase& multiphase)
{
	// The index of the given species in the (multiphase) instance
	const Index i = multiphase[species];
	
	// Define the right-hand side function h = h(T, P, n, a) for this equilibrium condition
	FunctionTPna h = [=](double T, double P, const VectorXd& n, const VectorXd& a)
		{ return n[i]/a[i] * activity; };
	
	return SpeciesMassConditionHelper2(species, h, multiphase);
}

const EquilibriumCondition ElementMassCondition2(const string& element, double mass, const Multiphase& multiphase)
{
	// The species in the system
	const vector<string>& systemSpecies = multiphase.GetSpecies();
	 
	// The number of species in the system
	const unsigned Ni = multiphase.GetNumSpecies();
	
	// The equilibrium condition instance
	EquilibriumCondition eCondition(Ni);
	
	// Set the alphaRow entries of (eCondition) with the number of atoms of (element) of each species
	for(unsigned i = 0; i < Ni; ++i) 
		eCondition.alpha_row[i] = ElementAtoms(systemSpecies[i], element);
	
	// Set the equilibrium condition function (h)
	eCondition.hj = ConstantTPna(mass);
	
	return eCondition;
}

const EquilibriumCondition AcidityCondition2(double pH, const Multiphase& multiphase)
{
	return SpeciesActivityCondition2("H[+]", pow(10.0, -pH), multiphase);
}

const EquilibriumCondition PartialPressureCondition2(const string& gas, double Pg, const Multiphase& multiphase)
{
	// Find the gaseous phase among all the phases in the system
	const Phase& gaseousPhase = multiphase.GetPhase("Gaseous");
	
	// The indexes of all gaseous species
	vector<Index> igases = multiphase[gaseousPhase.GetSpecies()];
	
	// Define the right-hand side function h = h(T, P, n, a) for this equilibrium condition
	auto h = [=](double T, double P, const VectorXd& n, const VectorXd& a)
	{ 
		// The total number of moles of gaseous species
		double nG = 0.0; for(unsigned i = 0; i < igases.size(); ++i) nG += n[igases[i]];
		
		return nG * Pg/P;
	};
	
	return SpeciesMassConditionHelper2(gas, h, multiphase);
}

const EquilibriumCondition ChargeBalanceCondition2(const Multiphase& multiphase)
{
	// The species in the system
	const vector<string>& systemSpecies = multiphase.GetSpecies();
	 
	// The number of species in the system
	const unsigned Ni = multiphase.GetNumSpecies();
	
	// The equilibrium condition instance
	EquilibriumCondition eCondition(Ni);
	
	// Set the alphaRow entries of (eCondition) with the electrical charge of each species
	for(unsigned i = 0; i < Ni; ++i) 
		eCondition.alpha_row[i] = ElectricalCharge(systemSpecies[i]);
	
	// Set the equilibrium condition function (h)
	eCondition.hj = ConstantTPna(0.0);
	
	return eCondition;
}
